import LSTM_Attention as model
import tensorflow as tf
import numpy as np

class PoseEstimate:
    def __init__(self, latest_ckpt):
        self.x = tf.placeholder(tf.float32, [None, 10, 68])
        self.predict = model.model(self.x)
        self.softmax_predict = tf.nn.softmax(self.predict)
        # self.graph = tf.Graph()
        self.sess = tf.Session()
        self.sess.run(tf.global_variables_initializer())

        saver = tf.train.Saver(tf.global_variables())
        model.load_model(self.sess, saver, latest_ckpt)
        # with self.graph.as_default():
        # latest_ckpt = tf.train.latest_checkpoint(latest_ckpt)
        # saver = tf.train.import_meta_graph(latest_ckpt + '.meta',
        #                                    clear_devices=True)
        # saver = tf.train.Saver(tf.global_variables())
        # saver.restore(self.sess, latest_ckpt)

    def estimate(self, data):
        # image = np.array(data)
        # image = image.reshape(1, 224, 224, 3)
        softmax_predict = self.sess.run([self.softmax_predict],
                                       feed_dict={self.x: data})
        return softmax_predict

if __name__ == '__main__':
    pe = PoseEstimate("./lstm_attention-all-1/model")
    data = np.random.random([1, 10, 68])
    # print(data)
    pre = pe.estimate(data)
    print(pre[0][0][1])